"""Example for the Euler equations"""

from hogs.grids import grid1d as grid
from hogs.solvers.euler1d import EESolver, EELagrangianSolver
from hogs.solvers.flux.euler.eulerian_flux import GodunovFluxEulerian, \
     EulerLxf, WAFFluxEulerian

from hogs.solvers.flux.euler.lagrangian_flux import LxfFluxLagrangian,\
     GodunovFluxLagrangian

from hogs.solvers.time_step_functions import EulerEquationsTimeStep
from hogs.solvers.primitive_variable_functions import EEPrimitiveVariable,\
     EEPrimitiveVariableLagrangian

import numpy

# parameters
gamma = 1.4

# create and initialize the grid
g = grid.Grid1D()
g.initialize(xlow=-1.0, xhigh=1.0, dx=0.005, nb=2, nvar=3)

# construct a one dimensional solver
solver = EELagrangianSolver(gamma=gamma, tf=0.6, nvar=3, grid=g)
#solver = EESolver(gamma=gamma, tf=0.1, nvar=3)

# set the flux function
#solver.flux_function = flux.GodunovFluxEulerian(gamma=1.4)
solver.flux_function = GodunovFluxLagrangian()
#solver.flux_function = LxfFluxLagrangian()
#solver.flux_function = EulerLxf()

# set the grid for the flux function
solver.flux_function.set_grid( solver.grid )

# set the time step function
solver.time_step_function = EulerEquationsTimeStep(grid=solver.grid)

# primitive variable function
solver.primitive_variable_function = EEPrimitiveVariable(gamma=gamma,
                                                         grid=solver.grid)

# process command line 
solver.setup()

# set the variables
grid = solver.grid

# cell centers
x = grid.xc
q = grid.q

# initial data for the shock tube problem
rhol = 1.0; rhor = 1.0
el = er = 1e-6
pl = pr = (gamma-1.0) * rhol * el
ul = 1.0; ur = -1.0

for i, j in enumerate(x):
    if j < 0.0:

        q[0,i] = rhol
        q[1,i] = rhol*ul

        q[2,i] = rhol * (0.5 * ul**2 + el)

    else:

        q[0,i] = rhor
        q[1,i] = rhor*ur

        q[2,i] = rhor * (0.5 * ur**2 + er)

solver.solve()

# from hogs.solvers.riemann import reuler
# # compute the exact solution
# rsolver = reuler.RiemannSolverEulerExact(gamma=1.4)

# pm, um = rsolver.compute_star(rhol, rhor, pl, pr, ul, ur)

# x = solver.grid.xc
# rho = numpy.ones_like(x)
# u = numpy.ones_like(x)
# p = numpy.ones_like(x)

# dx = x[1] - x[0]
# ncells = solver.grid.ncells

# for i in range(ncells):
#     s = (x[i] - 0.6)/0.15

#     _rho, _u, _p = rsolver.sample(pm, um, s, rhol, rhor, pl, pr, ul, ur)

#     rho[i] = _rho; u[i] = _u; p[i] = _p

# numpy.savez('exact', x=x, p=p, rho=rho, u=u)
